import os

from flask import Flask, redirect, url_for, Response, render_template, request, flash
from imageai.Classification import ImageClassification
from imageai.Detection import ObjectDetection
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['SECRET_KEY'] = "wb04307201-image-ai"
BASE_DIR = os.path.abspath(os.path.dirname(__file__))
UPLOAD_FOLDER = os.path.join(BASE_DIR, 'uploads')
RESULT_FOLDER = os.path.join(BASE_DIR, 'result')
MODELS_FOLDER = os.path.join(BASE_DIR, 'models')


@app.route('/')
def reindex():
    return redirect(url_for('index'))


@app.route('/index')
def index():
    return render_template('index.html')


@app.route('/ImagePrediction', methods=['POST', 'GET'])
def image_prediction():
    if request.method == 'POST':
        if 'file' not in request.files or request.files['file'].filename == '':
            flash('请先选择文件')
            return redirect(request.url)
        else:
            file = request.files['file']
            models = request.form['models']
            uploadname = secure_filename(file.filename)
            make_dir(UPLOAD_FOLDER)
            path = os.path.join(UPLOAD_FOLDER, uploadname)
            if os.path.exists(path):
                os.remove(path)
            file.save(path)

        prediction = ImageClassification()
        if models == 'mobilenet_v2-b0353104.pth':
            prediction.setModelTypeAsMobileNetV2()
        elif models == 'inception_v3_google-1a9a5a14.pth':
            prediction.setModelTypeAsInceptionV3()
        elif models == 'densenet121-a639ec97.pth':
            prediction.setModelTypeAsDenseNet121()
        else:
            prediction.setModelTypeAsResNet50()

        model_path = os.path.join(MODELS_FOLDER, models)
        check = os.path.exists(model_path)
        if not check:
            raise Exception("model is out of range")
        prediction.setModelPath(model_path)
        prediction.loadModel()

        predictions, probabilities = prediction.classifyImage(path, result_count=5)
        res_list = []
        for eachPrediction, eachProbability in zip(predictions, probabilities):
            res_list.append({'eachPrediction': eachPrediction, 'eachProbability': eachProbability})
            # print(eachPrediction + " : " + eachProbability)
        return render_template('ImagePrediction.html', filename=uploadname, resList=res_list)
    return render_template('ImagePrediction.html')


@app.route('/ObjectDetection', methods=['POST', 'GET'])
def object_detection():
    if request.method == 'POST':
        if 'file' not in request.files or request.files['file'].filename == '':
            flash('请先选择文件')
            return redirect(request.url)
        else:
            file = request.files['file']
            models = request.form['models']
            speed = request.form['speed']
            uploadname = secure_filename(file.filename)
            make_dir(UPLOAD_FOLDER)
            make_dir(RESULT_FOLDER)
            path = os.path.join(UPLOAD_FOLDER, uploadname)
            if os.path.exists(path):
                os.remove(path)
            respath = os.path.join(RESULT_FOLDER, uploadname)
            if os.path.exists(respath):
                os.remove(respath)
            file.save(path)

            detector = ObjectDetection()
            if models == 'mobilenet_v2-b0353104.pth':
                detector.setModelTypeAsRetinaNet()
            elif models == 'tiny-yolov3.pt':
                detector.setModelTypeAsTinyYOLOv3()
            else:
                detector.setModelTypeAsYOLOv3()

            model_path = os.path.join(MODELS_FOLDER, models)
            check = os.path.exists(model_path)
            if not check:
                raise Exception("model is out of range")
            detector.setModelPath(model_path)
            detector.loadModel()

            detections = detector.detectObjectsFromImage(input_image=path,
                                                         output_image_path=respath,
                                                         minimum_percentage_probability=30)
            # for eachObject in detections:
            #             #     print(eachObject["name"] + " : " + eachObject["percentage_probability"])
            return render_template('ObjectDetection.html', filename=uploadname)
    return render_template('ObjectDetection.html')


def make_dir(make_dir_path):
    path = make_dir_path.strip()
    if not os.path.exists(path):
        os.makedirs(path)


@app.route('/image/<type>/<filename>/')
def image(type, filename=None):
    if filename:
        if type == 'uploads':
            path = os.path.join(UPLOAD_FOLDER, filename)
        else:
            path = os.path.join(RESULT_FOLDER, filename)
        file = open(path, 'rb')
        return Response(file.read(), mimetype='image/jpeg')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port='5000')
